#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GPU服务器自动注册脚本
GPU服务器启动时自动注册到API网关
"""

import requests
import os
import time
import socket
import logging
from typing import Tuple, Optional

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def get_current_ip() -> str:
    """获取当前服务器的公网IP地址"""
    try:
        # 方法1：通过连接外部服务器获取本机IP
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        # 方法2：如果上述方法失败，可以调用外部API获取公网IP
        try:
            response = requests.get('https://api.ipify.org', timeout=5)
            return response.text.strip()
        except Exception:
            return "127.0.0.1"  # 最后的回退选项

def register_to_gateway(gateway_url: str = None, port: int = None, secret: str = None) -> Tuple[bool, Optional[str]]:
    """GPU服务器启动时注册到API网关"""
    
    # 获取配置参数
    gateway_url = gateway_url or os.getenv('GATEWAY_URL', 'https://gongjuxiang.work')
    port = port or int(os.getenv('PORT', '8000'))
    secret = secret or os.getenv('SHARED_SECRET', 'gpu-server-register-to-api-gateway-2024')
    
    # 获取当前IP地址
    ip = get_current_ip()
    
    # 构建注册URL
    register_url = f"{gateway_url}/webhook/register"
    
    # 准备注册数据
    data = {
        "ip": ip,
        "port": port,
        "secret": secret
    }
    
    logger.info(f"正在注册到API网关: {register_url}")
    logger.info(f"注册参数: IP={ip}, Port={port}")
    
    try:
        response = requests.post(register_url, json=data, timeout=10)
        
        if response.status_code == 200:
            result = response.json()
            if result.get('success'):
                assigned_server_id = result.get('server_id')
                logger.info(f"✅ 服务器 {assigned_server_id} 注册成功")
                return True, assigned_server_id
            else:
                logger.error(f"❌ 注册失败: {result.get('error')}")
                return False, None
        else:
            logger.error(f"❌ 注册请求失败: HTTP {response.status_code}")
            return False, None
            
    except Exception as e:
        logger.error(f"❌ 注册异常: {e}")
        return False, None

def unregister_from_gateway(gateway_url: str = None, port: int = None, secret: str = None) -> bool:
    """GPU服务器关闭时从API网关注销"""
    
    # 获取配置参数
    gateway_url = gateway_url or os.getenv('GATEWAY_URL', 'https://gongjuxiang.work')
    port = port or int(os.getenv('PORT', '8000'))
    secret = secret or os.getenv('SHARED_SECRET', 'gpu-server-register-to-api-gateway-2024')
    
    # 获取当前IP地址
    ip = get_current_ip()
    
    # 构建注销URL
    unregister_url = f"{gateway_url}/webhook/unregister"
    
    # 准备注销数据
    data = {
        "ip": ip,
        "port": port,
        "secret": secret
    }
    
    logger.info(f"正在从API网关注销: {unregister_url}")
    logger.info(f"注销参数: IP={ip}, Port={port}")
    
    try:
        response = requests.post(unregister_url, json=data, timeout=10)
        
        if response.status_code == 200:
            result = response.json()
            if result.get('success'):
                logger.info("✅ 服务器注销成功")
                return True
            else:
                logger.error(f"❌ 注销失败: {result.get('error')}")
                return False
        else:
            logger.error(f"❌ 注销请求失败: HTTP {response.status_code}")
            return False
            
    except Exception as e:
        logger.error(f"❌ 注销异常: {e}")
        return False

def main():
    """主函数 - 自动注册"""
    logger.info("🚀 GPU服务器自动注册启动")
    
    # 注册到API网关
    success, server_id = register_to_gateway()
    
    if success:
        logger.info(f"🎉 注册成功，服务器ID: {server_id}")
        
        # 保存服务器ID到环境变量或文件
        os.environ['SERVER_ID'] = server_id
        
        # 定期心跳检查（可选）
        logger.info("💓 开始心跳检查...")
        while True:
            time.sleep(30)  # 每30秒检查一次
            # 这里可以添加健康检查逻辑
            
    else:
        logger.error("💥 注册失败，退出程序")
        return 1
    
    return 0

if __name__ == '__main__':
    exit(main())
